import os
import torch
# 引入模型
from transformers import BertTokenizer, BertForSequenceClassification, pipeline
# 本地离线调用
from transformers import AutoModelForSequenceClassification, AutoTokenizer


# 远程下载模型
def download_model():
    """
    从远程下载预训练的 BERT 模型和分词器。

    Returns:
        base_model: 下载的 BERT 模型。
        base_tokenizer: 下载的 BERT 分词器。
    """
    # 设置具体包含config.json的目录，只支持绝对路径
    model_dir = r"model\bert-base-chinese"
    base_model = BertForSequenceClassification.from_pretrained("bert-base-chinese", cache_dir=model_dir)
    base_tokenizer = BertTokenizer.from_pretrained("bert-base-chinese", cache_dir=model_dir)
    return base_model, base_tokenizer


# 本地离线使用模型
def use_model():
    """
    从本地加载预训练的 BERT 模型和分词器。

    Returns:
        base_model: 本地加载的 BERT 模型。
        base_tokenizer: 本地加载的 BERT 分词器。
    """
    path = os.path.dirname(os.path.abspath(__file__))
    snapshot_dir = os.path.join(path, "model/bert-base-chinese/models--bert-base-chinese/snapshots")
    versions = os.listdir(snapshot_dir)
    latest_version = sorted(versions)[-1]
    model_dir = os.path.join(snapshot_dir, latest_version)
    # 使用本地模型路径加载模型和分词器，避免重复下载
    base_model = AutoModelForSequenceClassification.from_pretrained(model_dir)
    base_tokenizer = AutoTokenizer.from_pretrained(model_dir)
    return base_model, base_tokenizer


if __name__ == '__main__':
    # 设置设备
    device = "cuda" if torch.cuda.is_available() else "cpu"
    # model, tokenizer = download_model()
    model, tokenizer = use_model()
    print(model)

    # 创建分类 pipeline
    classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, device=device)

    # 进行文本分类
    result = classifier("我今天很难过")
    print(result)



